#!/usr/bin/env python3
import argparse
import os
import sys
from time import sleep

import grpc

# Import P4Runtime lib from parent utils dir
sys.path.append(
    os.path.join(os.path.dirname(os.path.abspath(__file__)),
                 '../../utils/'))
import p4runtime_lib.bmv2
import p4runtime_lib.helper
from p4runtime_lib.error_utils import printGrpcError
from p4runtime_lib.switch import ShutdownAllSwitchConnections

SWITCH_TO_HOST_PORT = 1

def writeTunnelRules(p4info_helper, ingress_sw, egress_sw, tunnel_id,
                     dst_eth_addr, dst_ip_addr):
    """
    配置隧道规则：入口规则、转发规则和出口规则
    """
    # 端口映射配置
    if tunnel_id == 100:  # s1 -> s2
        SWITCH_TO_SWITCH_PORT = 2
    elif tunnel_id == 101:  # s2 -> s1
        SWITCH_TO_SWITCH_PORT = 2
    elif tunnel_id == 200:  # s1 -> s3
        SWITCH_TO_SWITCH_PORT = 3
    elif tunnel_id == 201:  # s3 -> s1
        SWITCH_TO_SWITCH_PORT = 2
    elif tunnel_id == 300:  # s2 -> s3
        SWITCH_TO_SWITCH_PORT = 3
    elif tunnel_id == 301:  # s3 -> s2
        SWITCH_TO_SWITCH_PORT = 3
    else:
        print(f"Error: Unknown tunnel ID {tunnel_id}")
        return

    print(f"Configuring tunnel {tunnel_id}: {ingress_sw.name} -> {egress_sw.name}")

    # 1) Tunnel Ingress Rule - 只在入口交换机上配置
    table_entry = p4info_helper.buildTableEntry(
        table_name="MyIngress.ipv4_lpm",
        match_fields={
            "hdr.ipv4.dstAddr": (dst_ip_addr, 32)
        },
        action_name="MyIngress.myTunnel_ingress",
        action_params={
            "dst_id": tunnel_id,
        })
    ingress_sw.WriteTableEntry(table_entry)
    print(f"  Added ingress rule on {ingress_sw.name}: {dst_ip_addr} -> tunnel {tunnel_id}")

    # 2) Tunnel Transit Rule on Ingress Switch - 转发到下一跳
    table_entry = p4info_helper.buildTableEntry(
        table_name="MyIngress.myTunnel_exact",
        match_fields={
            "hdr.myTunnel.dst_id": tunnel_id
        },
        action_name="MyIngress.myTunnel_forward",
        action_params={
            "port": SWITCH_TO_SWITCH_PORT
        })
    ingress_sw.WriteTableEntry(table_entry)
    print(f"  Added forward rule on {ingress_sw.name}: tunnel {tunnel_id} -> port {SWITCH_TO_SWITCH_PORT}")

    # 3) Tunnel Egress Rule - 只在出口交换机上配置
    table_entry = p4info_helper.buildTableEntry(
        table_name="MyIngress.myTunnel_exact",
        match_fields={
            "hdr.myTunnel.dst_id": tunnel_id
        },
        action_name="MyIngress.myTunnel_egress",
        action_params={
            "dstAddr": dst_eth_addr,
            "port": SWITCH_TO_HOST_PORT
        })
    egress_sw.WriteTableEntry(table_entry)
    print(f"  Added egress rule on {egress_sw.name}: tunnel {tunnel_id} -> {dst_eth_addr}")

def writeArpRules(p4info_helper, sw, target_ip, mac_addr):
    """
    配置ARP回复规则 - 修正的参数格式
    """
    try:
        table_entry = p4info_helper.buildTableEntry(
            table_name="MyIngress.arp_match",
            match_fields={
                "hdr.arp.oper": 1,  # ARP请求操作码
                "hdr.arp.tpa": (target_ip, 32)  # 目标协议地址
            },
            action_name="MyIngress.send_arp_reply",
            action_params={
                "macAddr": mac_addr
            })
        sw.WriteTableEntry(table_entry)
        print(f"Added ARP rule on {sw.name}: {target_ip} -> {mac_addr}")
    except Exception as e:
        print(f"Error adding ARP rule on {sw.name}: {e}")

def writeDirectArpRules(p4info_helper, sw, host_ips, host_macs):
    """
    为直连主机配置ARP回复规则
    """
    for ip, mac in zip(host_ips, host_macs):
        try:
            table_entry = p4info_helper.buildTableEntry(
                table_name="MyIngress.arp_match",
                match_fields={
                    "hdr.arp.oper": 1,
                    "hdr.arp.tpa": (ip, 32)
                },
                action_name="MyIngress.send_arp_reply",
                action_params={
                    "macAddr": mac
                })
            sw.WriteTableEntry(table_entry)
            print(f"  Added direct ARP rule on {sw.name}: {ip} -> {mac}")
        except Exception as e:
            print(f"  Error adding direct ARP rule: {e}")

def readTableRules(p4info_helper, sw):
    """
    Reads the table entries from all tables on the switch.
    """
    print('\n----- Reading tables rules for %s -----' % sw.name)
    try:
        for response in sw.ReadTableEntries():
            for entity in response.entities:
                entry = entity.table_entry
                table_name = p4info_helper.get_tables_name(entry.table_id)
                print('%s: ' % table_name, end=' ')
                for m in entry.match:
                    print(p4info_helper.get_match_field_name(table_name, m.field_id), end=' ')
                    print('%r' % (p4info_helper.get_match_field_value(m),), end=' ')
                action = entry.action.action
                action_name = p4info_helper.get_actions_name(action.action_id)
                print('->', action_name, end=' ')
                for p in action.params:
                    print(p4info_helper.get_action_param_name(action_name, p.param_id), end=' ')
                    print('%r' % p.value, end=' ')
                print()
    except Exception as e:
        print(f"Error reading table rules: {e}")

def printtxt(p4info_helper, insw, esw, counter_name, index):
    """
    将计数器数据写入文件
    """
    try:
        for response in insw.ReadCounters(p4info_helper.get_counters_id(counter_name), index):
            for entity in response.entities:
                counter = entity.counter_entry
                filename = f"{insw.name}{esw.name}.txt"
                with open(filename, "a") as file:
                    file.write("%s->%s\n" % (insw.name, esw.name))
                    file.write("%s sent %d packets, counter ID: %d\n" % (insw.name, counter.data.packet_count, index))

        for response in esw.ReadCounters(p4info_helper.get_counters_id(counter_name), index):
            for entity in response.entities:
                counter = entity.counter_entry
                filename = f"{insw.name}{esw.name}.txt"
                with open(filename, "a") as file:
                    file.write("%s received %d packets, counter ID: %d\n" % (esw.name, counter.data.packet_count, index))
    except Exception as e:
        print(f"Error writing to file: {e}")

def printCounter(p4info_helper, sw, counter_name, index):
    """
    Reads the specified counter at the specified index from the switch.
    """
    try:
        for response in sw.ReadCounters(p4info_helper.get_counters_id(counter_name), index):
            for entity in response.entities:
                counter = entity.counter_entry
                print("%s %s %d: %d packets (%d bytes)" % (
                    sw.name, counter_name, index,
                    counter.data.packet_count, counter.data.byte_count
                ))
    except Exception as e:
        print(f"Error reading counter: {e}")

def main(p4info_file_path, bmv2_file_path):
    p4info_helper = p4runtime_lib.helper.P4InfoHelper(p4info_file_path)

    try:
        # Create switch connections
        print("Connecting to switches...")
        s1 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name='s1', address='127.0.0.1:50051', device_id=0,
            proto_dump_file='logs/s1-p4runtime-requests.txt')
        s2 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name='s2', address='127.0.0.1:50052', device_id=1,
            proto_dump_file='logs/s2-p4runtime-requests.txt')
        s3 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name='s3', address='127.0.0.1:50053', device_id=2,
            proto_dump_file='logs/s3-p4runtime-requests.txt')

        # Establish master arbitration
        print("Establishing controller as master...")
        s1.MasterArbitrationUpdate()
        s2.MasterArbitrationUpdate()
        s3.MasterArbitrationUpdate()

        # Install P4 program
        print("Installing P4 program on switches...")
        s1.SetForwardingPipelineConfig(p4info=p4info_helper.p4info,
                                       bmv2_json_file_path=bmv2_file_path)
        s2.SetForwardingPipelineConfig(p4info=p4info_helper.p4info,
                                       bmv2_json_file_path=bmv2_file_path)
        s3.SetForwardingPipelineConfig(p4info=p4info_helper.p4info,
                                       bmv2_json_file_path=bmv2_file_path)

        # Write tunnel rules to s1, s2, s3
        print("\nConfiguring tunnel rules...")

        # s1 -> s2 (h1 -> h2)
        writeTunnelRules(p4info_helper, ingress_sw=s1, egress_sw=s2, tunnel_id=100,
                        dst_eth_addr="08:00:00:00:02:22", dst_ip_addr="10.0.2.2")

        # s1 -> s3 (h1 -> h3)
        writeTunnelRules(p4info_helper, ingress_sw=s1, egress_sw=s3, tunnel_id=200,
                        dst_eth_addr="08:00:00:00:03:33", dst_ip_addr="10.0.3.3")

        # s2 -> s1 (h2 -> h1)
        writeTunnelRules(p4info_helper, ingress_sw=s2, egress_sw=s1, tunnel_id=101,
                        dst_eth_addr="08:00:00:00:01:11", dst_ip_addr="10.0.1.1")

        # s2 -> s3 (h2 -> h3)
        writeTunnelRules(p4info_helper, ingress_sw=s2, egress_sw=s3, tunnel_id=300,
                        dst_eth_addr="08:00:00:00:03:33", dst_ip_addr="10.0.3.3")

        # s3 -> s1 (h3 -> h1)
        writeTunnelRules(p4info_helper, ingress_sw=s3, egress_sw=s1, tunnel_id=201,
                        dst_eth_addr="08:00:00:00:01:11", dst_ip_addr="10.0.1.1")

        # s3 -> s2 (h3 -> h2)
        writeTunnelRules(p4info_helper, ingress_sw=s3, egress_sw=s2, tunnel_id=301,
                        dst_eth_addr="08:00:00:00:02:22", dst_ip_addr="10.0.2.2")

        # Write arp_match table entries to s1, s2, s3 - 修正的ARP配置
        print("\nConfiguring ARP rules...")

        # 配置交换机接口的ARP回复
        writeArpRules(p4info_helper, s1, "10.0.1.10", "08:00:00:00:01:00")
        writeArpRules(p4info_helper, s2, "10.0.2.20", "08:00:00:00:02:00")
        writeArpRules(p4info_helper, s3, "10.0.3.30", "08:00:00:00:03:00")

        # 为直连主机配置ARP回复
        print("Configuring direct host ARP rules...")
        writeDirectArpRules(p4info_helper, s1, ["10.0.1.1"], ["08:00:00:00:01:11"])
        writeDirectArpRules(p4info_helper, s2, ["10.0.2.2"], ["08:00:00:00:02:22"])
        writeDirectArpRules(p4info_helper, s3, ["10.0.3.3"], ["08:00:00:00:03:33"])

        # Read table entries from s1, s2 and s3
        print("\nReading table entries...")
        readTableRules(p4info_helper, s1)
        readTableRules(p4info_helper, s2)
        readTableRules(p4info_helper, s3)

        print("\n" + "="*60)
        print("Configuration completed successfully!")
        print("Host connectivity map:")
        print("  h1 (10.0.1.1) <-> h2 (10.0.2.2) via tunnel 100/101")
        print("  h1 (10.0.1.1) <-> h3 (10.0.3.3) via tunnel 200/201")
        print("  h2 (10.0.2.2) <-> h3 (10.0.3.3) via tunnel 300/301")
        print("="*60)

        # 清理之前的计数器文件
        for filename in ["s1s2.txt", "s2s1.txt", "s1s3.txt", "s3s1.txt", "s2s3.txt", "s3s2.txt"]:
            if os.path.exists(filename):
                os.remove(filename)

        # Print the tunnel counters every 2 seconds
        print("\nStarting counter monitoring (Ctrl+C to stop)...")
        counter = 0
        while True:
            sleep(2)
            counter += 1
            print(f'\n----- Reading tunnel counters (iteration {counter}) -----')

            print('\n----- s1 -> s2 -----')
            printCounter(p4info_helper, s1, "MyIngress.ingressTunnelCounter", 100)
            printCounter(p4info_helper, s2, "MyIngress.egressTunnelCounter", 100)
            printtxt(p4info_helper, s1, s2, "MyIngress.ingressTunnelCounter", 100)

            print('\n----- s2 -> s1 -----')
            printCounter(p4info_helper, s2, "MyIngress.ingressTunnelCounter", 101)
            printCounter(p4info_helper, s1, "MyIngress.egressTunnelCounter", 101)
            printtxt(p4info_helper, s2, s1, "MyIngress.ingressTunnelCounter", 101)

            print('\n----- s1 -> s3 -----')
            printCounter(p4info_helper, s1, "MyIngress.ingressTunnelCounter", 200)
            printCounter(p4info_helper, s3, "MyIngress.egressTunnelCounter", 200)
            printtxt(p4info_helper, s1, s3, "MyIngress.ingressTunnelCounter", 200)

            print('\n----- s3 -> s1 -----')
            printCounter(p4info_helper, s3, "MyIngress.ingressTunnelCounter", 201)
            printCounter(p4info_helper, s1, "MyIngress.egressTunnelCounter", 201)
            printtxt(p4info_helper, s3, s1, "MyIngress.ingressTunnelCounter", 201)

            print('\n----- s2 -> s3 -----')
            printCounter(p4info_helper, s2, "MyIngress.ingressTunnelCounter", 300)
            printCounter(p4info_helper, s3, "MyIngress.egressTunnelCounter", 300)
            printtxt(p4info_helper, s2, s3, "MyIngress.ingressTunnelCounter", 300)

            print('\n----- s3 -> s2 -----')
            printCounter(p4info_helper, s3, "MyIngress.ingressTunnelCounter", 301)
            printCounter(p4info_helper, s2, "MyIngress.egressTunnelCounter", 301)
            printtxt(p4info_helper, s3, s2, "MyIngress.ingressTunnelCounter", 301)

            print('\n------------ Finished ------------')

    except KeyboardInterrupt:
        print("\nShutting down controller.")
    except grpc.RpcError as e:
        printGrpcError(e)
    except Exception as e:
        print(f"Unexpected error: {e}")

    ShutdownAllSwitchConnections()

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='P4Runtime Controller')
    parser.add_argument('--p4info', help='p4info proto in text format from p4c',
                        type=str, action="store", required=False,
                        default='./build/advanced_tunnel.p4.p4info.txt')
    parser.add_argument('--bmv2-json', help='BMv2 JSON file from p4c',
                        type=str, action="store", required=False,
                        default='./build/advanced_tunnel.json')
    args = parser.parse_args()

    if not os.path.exists(args.p4info):
        parser.print_help()
        print("\np4info file not found: %s\nHave you run 'make'?" % args.p4info)
        parser.exit(1)
    if not os.path.exists(args.bmv2_json):
        parser.print_help()
        print("\nBMv2 JSON file not found: %s\nHave you run 'make'?" % args.bmv2_json)
        parser.exit(1)

    # Create logs directory if it doesn't exist
    os.makedirs('logs', exist_ok=True)

    main(args.p4info, args.bmv2_json)
